#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
文档质量优化数据集向量化脚本

对已有的标签好的数据进行embedding处理，并按标签分类存储。
使用sentence-transformers构建向量库，支持快速相似性检索。

"""

import json
import os
import pickle
import time
from collections import defaultdict
from pathlib import Path
from typing import Dict, List, Tuple, Any

import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from tqdm import tqdm
from config import Config


class DocumentEmbeddingBuilder:
    """文档嵌入向量构建器"""
    
    def __init__(self, 
                 model_name: str = None,
                 output_dir: str = None):
        """
        初始化构建器
        
        Args:
            model_name: SentenceTransformers模型名称，选择支持中文的多语言模型
            output_dir: 输出目录
        """
        self.model_name = model_name or Config.get_embedding_model()
        self.output_dir = Path(output_dir or Config.get_embedding_dir())
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        print(f"正在加载模型: {self.model_name}")
        self.model = SentenceTransformer(self.model_name)
        print(f"模型加载完成，embedding维度: {self.model.get_sentence_embedding_dimension()}")
        
    def extract_text_for_embedding(self, data_item: Dict[str, Any]) -> str:
        """
        提取用于embedding的文本内容
        
        将修改前后的内容、上下文合并为一个文本块
        
        Args:
            data_item: 单条数据项
            
        Returns:
            合并后的文本内容
        """
        # 提取关键字段
        add_content = data_item.get('add_content', '').strip()
        remove_content = data_item.get('remove_content', '').strip()
        context_before = data_item.get('context_before', '').strip()
        context_after = data_item.get('context_after', '').strip()
        
        # 构建组合文本，包含修改的前后内容和上下文
        text_parts = []
        
        if context_before:
            text_parts.append(f"修改前上下文: {context_before}")
            
        if remove_content:
            text_parts.append(f"修改前内容: {remove_content}")
            
        if add_content:
            text_parts.append(f"修改后内容: {add_content}")
            
        if context_after:
            text_parts.append(f"修改后上下文: {context_after}")
            
        # 合并所有文本
        combined_text = " ".join(text_parts)
        
        # 如果组合文本为空，使用文件路径作为备选
        if not combined_text.strip():
            combined_text = data_item.get('file_path', 'unknown_file')
            
        return combined_text
    
    def load_and_filter_data(self, data_file: str) -> Tuple[List[Dict], Dict[str, List]]:
        """
        加载并过滤有效的标签数据
        
        Args:
            data_file: 数据文件路径
            
        Returns:
            (有效数据列表, 按标签分组的数据)
        """
        print(f"正在加载数据文件: {data_file}")
        
        valid_data = []
        label_groups = defaultdict(list)
        
        line_count = 0
        valid_count = 0
        
        with open(data_file, 'r', encoding='utf-8') as f:
            for line in tqdm(f, desc="处理数据行"):
                line_count += 1
                try:
                    data = json.loads(line.strip())
                    evaluation = data.get('evaluation', {})
                    
                    # 只处理有效且有标签的数据
                    is_valid = evaluation.get('is_valid', False)
                    label_id = evaluation.get('label_id')
                    
                    if is_valid and label_id:
                        # 直接使用原始数据，只添加扁平化的标签信息
                        data['label_id'] = label_id
                        data['label_name'] = evaluation.get('label_name', '')
                        data['label_reason'] = evaluation.get('label_reason', '')
                        
                        valid_data.append(data)
                        label_groups[label_id].append(len(valid_data) - 1)  # 存储索引
                        valid_count += 1
                        
                except Exception as e:
                    print(f"解析第{line_count}行时出错: {e}")
                    continue
        
        print(f"数据加载完成:")
        print(f"  总行数: {line_count}")
        print(f"  有效样本数: {valid_count}")
        print(f"  标签类别数: {len(label_groups)}")
        
        # 输出标签分布统计
        print(f"\n标签分布:")
        for label_id, indices in sorted(label_groups.items()):
            label_name = valid_data[indices[0]]['label_name'] if indices else "未知"
            print(f"  标签{label_id} ({label_name}): {len(indices)}条")
            
        return valid_data, dict(label_groups)
    
    def build_embeddings(self, valid_data: List[Dict]) -> np.ndarray:
        """
        构建所有数据的embedding向量
        
        Args:
            valid_data: 有效数据列表
            
        Returns:
            embedding矩阵 (n_samples, embedding_dim)
        """
        print("正在提取文本内容...")
        texts = []
        for item in tqdm(valid_data, desc="提取文本"):
            text = self.extract_text_for_embedding(item)
            texts.append(text)
        
        print(f"正在计算{len(texts)}个文本的embedding...")
        
        # 批量计算embedding，避免内存问题
        batch_size = 32
        embeddings = []
        
        for i in tqdm(range(0, len(texts), batch_size), desc="计算embedding"):
            batch_texts = texts[i:i + batch_size]
            batch_embeddings = self.model.encode(
                batch_texts,
                convert_to_numpy=True,
                show_progress_bar=False
            )
            embeddings.append(batch_embeddings)
        
        # 合并所有批次的结果
        all_embeddings = np.vstack(embeddings)
        print(f"embedding计算完成，形状: {all_embeddings.shape}")
        
        return all_embeddings
    
    def build_faiss_index(self, embeddings: np.ndarray) -> faiss.IndexFlatIP:
        """
        构建FAISS索引以支持快速检索
        
        Args:
            embeddings: embedding矩阵
            
        Returns:
            FAISS索引
        """
        print("正在构建FAISS索引...")
        
        # 标准化向量以使用内积进行相似度计算
        faiss.normalize_L2(embeddings)
        
        # 使用内积索引（适合标准化后的向量）
        dimension = embeddings.shape[1]
        index = faiss.IndexFlatIP(dimension)
        index.add(embeddings)
        
        print(f"FAISS索引构建完成，包含{index.ntotal}个向量")
        return index
    
    def save_label_data(self, 
                       label_id: str,
                       label_data: List[Dict],
                       label_embeddings: np.ndarray,
                       label_index: faiss.IndexFlatIP):
        """
        保存单个标签的数据和向量库
        
        Args:
            label_id: 标签ID
            label_data: 该标签的数据列表
            label_embeddings: 该标签的embedding矩阵
            label_index: 该标签的FAISS索引
        """
        label_dir = self.output_dir / f"label_{label_id}"
        label_dir.mkdir(exist_ok=True)
        
        # 保存元数据
        metadata_file = label_dir / "metadata.json"
        metadata = {
            'label_id': label_id,
            'label_name': label_data[0]['label_name'] if label_data else '',
            'sample_count': len(label_data),
            'embedding_dimension': label_embeddings.shape[1],
            'model_name': self.model_name,
            'created_time': time.strftime('%Y-%m-%d %H:%M:%S')
        }
        
        with open(metadata_file, 'w', encoding='utf-8') as f:
            json.dump(metadata, f, ensure_ascii=False, indent=2)
        
        # 保存数据详情
        data_file = label_dir / "data.jsonl"
        with open(data_file, 'w', encoding='utf-8') as f:
            for item in label_data:
                f.write(json.dumps(item, ensure_ascii=False) + '\n')
        
        # 保存embedding矩阵
        embeddings_file = label_dir / "embeddings.npy"
        np.save(embeddings_file, label_embeddings)
        
        # 保存FAISS索引
        index_file = label_dir / "faiss_index.bin"
        faiss.write_index(label_index, str(index_file))
        
        print(f"标签{label_id}数据已保存到 {label_dir}")
    
    def run(self, data_file: str):
        """
        运行完整的向量化流程
        
        Args:
            data_file: 输入数据文件路径
        """
        print("=" * 60)
        print("开始构建文档质量优化数据集向量库")
        print("=" * 60)
        
        # 1. 加载和过滤数据
        valid_data, label_groups = self.load_and_filter_data(data_file)
        
        if not valid_data:
            print("没有找到有效数据，退出处理")
            return
        
        # 2. 计算所有数据的embedding
        all_embeddings = self.build_embeddings(valid_data)
        
        # 3. 为每个标签构建单独的向量库
        print(f"\n正在为{len(label_groups)}个标签构建向量库...")
        
        for label_id, indices in tqdm(label_groups.items(), desc="处理标签"):
            # 提取该标签的数据和embedding
            label_data = [valid_data[i] for i in indices]
            label_embeddings = all_embeddings[indices]
            
            # 构建该标签的FAISS索引
            label_index = self.build_faiss_index(label_embeddings.copy())
            
            # 保存该标签的向量库
            self.save_label_data(label_id, label_data, label_embeddings, label_index)
        
        # 4. 保存全局信息
        global_metadata = {
            'total_samples': len(valid_data),
            'total_labels': len(label_groups),
            'model_name': self.model_name,
            'embedding_dimension': all_embeddings.shape[1],
            'label_distribution': {label_id: len(indices) for label_id, indices in label_groups.items()},
            'created_time': time.strftime('%Y-%m-%d %H:%M:%S'),
            'data_source': data_file
        }
        
        global_metadata_file = self.output_dir / "global_metadata.json"
        with open(global_metadata_file, 'w', encoding='utf-8') as f:
            json.dump(global_metadata, f, ensure_ascii=False, indent=2)
        
        print(f"\n向量库构建完成!")
        print(f"输出目录: {self.output_dir}")
        print(f"总样本数: {len(valid_data)}")
        print(f"标签类别数: {len(label_groups)}")
        print("=" * 60)


def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='构建文档质量优化数据集向量库')
    parser.add_argument('--input-file', type=str, 
                       default='../data/label_data/auto_pass.jsonl',
                       help='输入数据文件路径')
    parser.add_argument('--output-dir', type=str,
                       default=Config.get_embedding_dir(),
                       help='输出目录路径')
    parser.add_argument('--model-name', type=str,
                       default=Config.get_embedding_model(),
                       help='SentenceTransformers模型名称')
    
    args = parser.parse_args()
    
    # 检查输入文件是否存在
    if not os.path.exists(args.input_file):
        print(f"错误: 数据文件不存在: {args.input_file}")
        return
    
    # 创建构建器并运行
    builder = DocumentEmbeddingBuilder(
        model_name=args.model_name,
        output_dir=args.output_dir
    )
    builder.run(args.input_file)


if __name__ == "__main__":
    main()
